################################################################################
# Author: Subhendu Malakar
# Mentor: Mridul Sankar Barik
# About: This code is to interact with the database for dragonsoft.com website.
#	The options are:
#	1. Create table : This is to create a new table if there
#	isn't any. Do not opt this option if there is an existing table.
#	2. Drop table : This option is to drop the table.
#	CAUTION: This option will delete all the data present in the database.
#	3. View table : This option is to view the tuples in the databse.
#	4. Drop and Create table: This option is to drop and create a new database.
################################################################################

# Importing psycopg2 module to connect with the PostGRESQL database.
import psycopg2

# For exit() and traceback for exceptions.
import sys,traceback

# This function creates a new database.
def create_table(conn,cur):
	
	# Enabling smooth operation using exceptions
	try:
		
		# execute(string) executes the string in the postgresql command.
		cur.execute("CREATE TABLE dragonsoft(cve_id varchar(11),dragonsoft_id integer NOT NULL,cvss_score float,description varchar(5000),attack_from varchar(15),impact varchar(100),links varchar(2000),products varchar(2000),solutions varchar(2000),PRIMARY KEY(cve_id))");
		
		# Committing all the changes.
		conn.commit();
		
	# Catching exceptions.
	except Exception as e: 
		print(str(e))

# This will drop the table.
def drop_database(cur):
	cur.execute("DROP TABLE dragonsoft")

# To view the database.
def view_database(cur):
	cur.execute("SELECT * FROM dragonsoft")

	# fetchall() fetches all the output from the previous execute() call.
	lst=cur.fetchall()
	
	# Printing all the tuples.
	for d in lst:
		print('\n\n\nCVE-ID: ',d[0])
		print('\nCVSS Score: ',d[2])
		print('\nDescription: ',d[3])
		print('\nAttack From: ',d[4])
		print('\nImpact: ',d[5])
		print('\nReferences: ',d[6])
		print('\nProducts Affected: ',d[7])
		print('\nSolutions: ',d[8])

# Main function. Execution starts here.
def main():
	
	# Creating a new connection to the database. This connection will be used for further interaction with the database.
	# For details on connect() see the report.
	try:
		conn=psycopg2.connect("dbname='postgres' user='postgres' host='localhost' port='5432' password='postgres'")
		
	# Catching exceptions.
	except:
		print('Database connection error.');
		
	# If everything goes fine.
	else:

		# Creating a cursor for the connection.
		cur=conn.cursor()
		
		# Taking input from the user for further course of action.
		inp=int(input('1.Create table\n2.Drop table\n3.View table\n4.Drop and create a table\n\n\n Enter your Choice(1/2/3/4):'))
		if inp==1:
			create_table(conn,cur)
		elif inp==2:
			drop_database(cur)
		elif inp==3:
			view_database(cur)
		elif inp==4:
			drop_database(cur)
			create_table(conn,cur)
		else:
			print('Wrong Input.\n Exiting...')
			
		# Closing cursor after the work is done.
		cur.close()
		
		# Committing connection.
		conn.commit()
		
	# Closing connection.
	conn.close();

# Ensuring main() is the first function to be called.
if __name__ == '__main__':
	main()
